feat(attention): fixed-m splash attention kernel with dynamic bounds and safety fallbacks - #477

…and safety fallbacks

Implements exact fixed-m splash attention in Pallas on TPU. Instead of an
online running max, a head / Q-block whose Cauchy-Schwarz bound
max||q|| * max||k - k_bar|| is within a length-dependent window uses a
precomputed m = ceil(bound) - C(N) (+ q . k_bar with virtual K-centering) and
skips the per-block rescales; everything else keeps online softmax.

- `get_fixed_m_constants(N)`: length-dependent recenter C(N) and safe bound
  that keep the f32 numerator/denominator from overflowing or flushing.
- `fixed_m_dtype_is_safe`: rejects fp16/fp8 (exponent range too small for the
  recentered exponent); bf16 and fp32 pass.
- V-magnitude gate: |V| <= DEFAULT_MAX_V_BOUND (256), otherwise online.
- Per-(head, Q-block) dispatch in the kernel; the Ulysses and R=1 ring callers
  compile a uniform (all-fixed) and a hybrid kernel and pick one with
  `lax.cond` on `all_fixed`. Both kernels are compiled; the extra compile
  time / program size has not been measured.
- Virtual K-centering: the kernel receives k_mean and adds q . k_bar to the
  fixed bound in VMEM (an MXU dot_general on an (Hkv, 8, d) tiled k_mean,
  precision=HIGHEST because the gate has no slack at its edge). K itself is
  never re-written in HBM. k_mean is padded to the lane-padded head_dim after
  `_pad_data_for_flash`, so head_dim < 128 works (regression test at d=64).

Explicit metadata contracts on the fixed-m ring path
----------------------------------------------------
1. Norm representation is declared, not inferred. The gate previously guessed
   whether `fixed_m_norms` were squared with `(qn.max() * mk.max()) < 1000.0`.
   Magnitude cannot answer that: legacy unsquared norms of (1000, 2) have a
   true bound of 2000, but read as squared give sqrt(2000) ~= 44.7, a ~45x
   under-estimate that admits fixed-m where it must fall back. Replaced by a
   REQUIRED `fixed_m_norms_squared` argument (no default; omitting it with
   fixed-m on raises). Tests use squared norms.

2. The V-safety predicate is required, not assumed. The V-magnitude and dtype
   verdict cannot be re-derived from one hop's Q/K, so the ring kernel raises
   unless `v_ok` is passed (`value=None` in the metadata helper gives
   v_ok = 0, fail-closed). `_ulysses_ring_custom_attention` (R > 1) computes it
   on the PRE-all-to-all shards -- dtype safety plus |V| <= 256 -- and reduces
   it with pmin over BOTH internal axes (ring, ulysses): before the a2a each
   device holds a sequence shard of every head, and the ring's
   accumulate-vs-LSE branch must be taken uniformly by every ppermute
   participant. Only Q and K go through the optimization_barrier that keeps
   XLA from duplicating their producer chains into the norm reductions; V is
   reduced outside it so the Q/K a2a does not wait on V's producer.

3. Norm shape is validated against per_q_block. Both gates compute
   `qn * mk[:, None]`, so a (num_heads,) array supplied with per_q_block=True
   broadcasts to (num_heads, num_heads) and pairs head j's query norm with head
   h's key norm; a sink head then inherits a small bound and the kernel
   evaluates exp2(large_logit - small_m) -> inf. The kernel now rejects the
   mismatch.

Known limitations
- On the uncentered R > 1 ring path the fixed bound can still flush the
  numerator for heads far below the bound; this is not measured here.
- The ring kernel runs its own pmin over global-bound eligibility (added by
  this commit) unless the caller promises `uniform_fixed_m=True`; on the
  Ulysses+Ring path that pmin is redundant with the caller's (not measured).

Tests
- custom_splash_fixed_m_test.py (31): kernel unit tests with relative-L2
  checks, k-norm (mk) shape and V-metadata fail-closed checks, integration
  through `_ulysses_attention` and the R=1 ring (uniform branch, hybrid
  branch via a sink head, d=64 -- the test re-derives the gate from the
  inputs and asserts it admits fixed-m, it does not observe the cond -- ragged
  KV tails with N % 8 == 0 and N % block_kv != 0 whose tail is an odd multiple
  of 8), and a U=2/R=2 check that a single V outlier on one pre-a2a device
  makes v_ok False on every device (ring kernel stubbed; needs >= 4 devices,
  explicit skip otherwise).
- ring_fixed_m_test.py (18): 8 backend-independent contract tests (including
  the required fixed_m_norms_squared and v_ok arguments), 6 ring numerical
  tests vs an f32 reference (2 devices), 1 test that declared-unsquared and
  squared norms give the same result, and 3 tests showing the raw-K bound is
  unsound where the centered bound is not.
- CI: to keep GitHub Actions short, 18 of these 49 tests (the kernel
  numerics grids and the multi-device integration tests) skip when
  GITHUB_ACTIONS=true. The 9 tests from main and the contract, dtype and
  V-metadata tests still run there. The new
  end_to_end/tpu/run_wan_stack_tests.sh runs all of them on a TPU VM; later
  PRs in this stack add their test files to it.

Verified (final tree): TPU v6e-8 (jax 0.11.2, end_to_end/tpu/run_wan_stack_tests.sh): 49 passed, 6 subtests passed in 170.4s.